package heartbeat

import (
	"log"
	"math/rand"
	"os"
	"strconv"
	"sync"
	"time"

	rabbitmq "dos_framework/heartbeat"
)

var dataServers = make(map[string]time.Time)
var mux sync.Mutex

func ListenHeartbeat() {
	q := rabbitmq.New(os.Getenv("RABBITMQ_SERVER"))
	defer q.Close()

	q.Bind("apiServers")
	c := q.Consume()
	go removeExpiredDataServer()
	for msg := range c {
		dataServer, e := strconv.Unquote(string(msg.Body))
		if e != nil {
			log.Panicf("unquote failed: %s", e)
		}
		mux.Lock()
		dataServers[dataServer] = time.Now()
		mux.Unlock()
	}

}
func removeExpiredDataServer() {

	for {
		time.Sleep(5 * time.Second)
		mux.Lock()

		for ds, t := range dataServers {
			if t.Add(10 * time.Second).Before(time.Now()) {
				delete(dataServers, ds)
			}
		}
		mux.Unlock()
	}
}
func GetDataServers() []string {
	mux.Lock()
	defer mux.Unlock()

	ds := make([]string, 0)
	for k, _ := range dataServers {
		ds = append(ds, k)
	}
	return ds
}

func ChooseRandomDataServer() string {
	ds := GetDataServers()
	if len(ds) == 0 {
		return ""
	}
	return ds[rand.Intn(len(ds))]

}

// 获取用于上传复原分片的随机数据服务节点， 排除掉已有分片所在的数据服务节点
// Parameters:
//
//	n:       Number of servers to select (must be > 0)
//	exclude: Map of server IDs to addresses that should not be selected
//
// Returns:
//
//	Slice containing network addresses of selected data servers (may be empty if no servers available)
func ChooseRandomDataServers(n int, exclude map[int]string) (ds []string) {
	candidates := make([]string, 0)
	reverseExcludeMap := make(map[string]int)

	for id, addr := range exclude {
		reverseExcludeMap[addr] = id
	}

	servers := GetDataServers()
	for i := range servers {
		if _, excluded := reverseExcludeMap[servers[i]]; !excluded {
			candidates = append(candidates, servers[i])
		}

	}
	length := len(candidates)
	if length < n {
		return
	}
	p := rand.Perm(length)
	for i := 0; i < n; i++ {
		ds = append(ds, candidates[p[i]])
	}
	return

}
